#include <iostream>
#include "TFile.h"
#include "TTree.h"
#include "TH1D.h"
#include "HelixHough.h"
#include "HelixRange.h"
#include "HelixResolution.h"
#include "VertexFitFunc.h"
#include "NewtonMinimizerGradHessian.h"
#include <sys/time.h>
#include <Eigen/LU>
#include <Eigen/Core>
#include <math.h>
#include "sPHENIXTracker.h"
#include "SimpleTrack.h"


using namespace std;
using namespace FitNewton;
using namespace Eigen;


int main(int argc, char** argv)
{
  bool print_truth = false;
  bool print_reco = false;
  
  unsigned int nthreads = 8;
  
  TFile infile(argv[1]);
  TTree* etree=0;
  infile.GetObject("events", etree);
  SimpleMCEvent* inputevent=0;
  etree->SetBranchAddress("tracks", &inputevent);
  
  int nlayers = 13;
  vector<float> radii;
  radii.assign(nlayers,0.);
  radii[0]=4.4;
  radii[1]=7.3;
  radii[2]=10.2;
  radii[3]=25.0;
  radii[4]=34.0;
  radii[5]=43.0;
  radii[6]=52.0;
  radii[7]=61.0;
  radii[8]=69.6;
  radii[9]=78.2;
  radii[10]=86.8;
  radii[11]=96.5;
  radii[12]=108.0;
  vector<float> material;
  material.assign(nlayers, 0.01);
  float Bfield = 4.;
  
  
  vector<float> smear_xy_layer;smear_xy_layer.assign(nlayers,0);
  vector<float> smear_z_layer;smear_z_layer.assign(nlayers,0);
  float sqrt_12 = sqrt(12.);
  for(int i=0;i<=2;++i)
  {
    smear_xy_layer[i] = 0.005;
    smear_z_layer[i] =  0.01;
  }
  for(int i=3;i<=4;++i)
  {
    smear_xy_layer[i] = 0.005;
    smear_z_layer[i] =  0.1;
  }
  for(int i=5;i<=6;++i)
  {
    smear_xy_layer[i] = 0.005;
    smear_z_layer[i] =  11.6;
  }
  for(int i=7;i<=8;++i)
  {
    smear_xy_layer[i] = 0.018;
    smear_z_layer[i] =  0.2;
  }
  for(int i=9;i<=12;++i)
  {
    smear_xy_layer[i] = 0.018;
    smear_z_layer[i] =  18.5;
  }
  
  float kappa_max = 0.03;
  float rho_max = pow(kappa_max,1.);
  
  
  //phi,d,kappa,dzdl,z0
  float min_iter = 2.;
  float max_iter = 9.;
  HelixRange top_range(0., 2.*M_PI,   -0.3, 0.3,   0.0, rho_max, -0.9, 0.9,   -0.15, 0.15);
  
  vector<unsigned int> onezoom(5,0);
  vector<vector<unsigned int> > zoomprofile;
  zoomprofile.assign(5,onezoom);
  for(unsigned int i=0;i<=0;++i)
  {
    zoomprofile[i][0] = 8;
    zoomprofile[i][1] = 1;
    zoomprofile[i][2] = 1;
    zoomprofile[i][3] = 1;
    zoomprofile[i][4] = 1;
  }
  for(unsigned int i=1;i<=1;++i)
  {
    zoomprofile[i][0] = 4;
    zoomprofile[i][1] = 1;
    zoomprofile[i][2] = 5;
    zoomprofile[i][3] = 6;
    zoomprofile[i][4] = 1;
  }
  for(unsigned int i=2;i<=2;++i)
  {
    zoomprofile[i][0] = 6;
    zoomprofile[i][1] = 5;
    zoomprofile[i][2] = 1;
    zoomprofile[i][3] = 12;
    zoomprofile[i][4] = 3;
  }
  for(unsigned int i=3;i<=3;++i)
  {
    zoomprofile[i][0] = 8;
    zoomprofile[i][1] = 1;
    zoomprofile[i][2] = 2;
    zoomprofile[i][3] = 2;
    zoomprofile[i][4] = 1;
  }
  for(unsigned int i=4;i<=4;++i)
  {
    zoomprofile[i][0] = 2;
    zoomprofile[i][1] = 1;
    zoomprofile[i][2] = 2;
    zoomprofile[i][3] = 2;
    zoomprofile[i][4] = 1;
  }
  sPHENIXTracker tracker(zoomprofile, 2, top_range, material, radii, Bfield, true, nthreads);
//   sPHENIXTracker tracker(zoomprofile, 2, top_range, material, radii, Bfield);
  tracker.setNLayers(5);
  unsigned int max_hits_seed = 20;
  unsigned int min_hits_seed = 5;
  tracker.setClusterStartBin(2);
  tracker.setRejectGhosts(true);
  tracker.setChi2Cut(6.0);
  tracker.setChi2RemovalCut(4.0);
  tracker.setPrintTimings(true);
  tracker.setVerbosity(1);
  tracker.setCutOnDca(false);
  tracker.requireLayers(5);
  tracker.setSmoothBack(false);
  tracker.setBinScale(0.6);
  tracker.setZBinScale(0.6);
  tracker.setRemoveHits(true);
  tracker.setSeparateByHelicity(true);
  tracker.setMaxHitsPairs(0);
  
  
  // *********************************************************
  // setup for tracking based on seeds
  vector<vector<unsigned int> > zoomprofile_2;
  zoomprofile_2.assign(7,onezoom);
  for(unsigned int i=0;i<=0;++i)
  {
    zoomprofile_2[i][0] = 8;
    zoomprofile_2[i][1] = 1;
    zoomprofile_2[i][2] = 1;
    zoomprofile_2[i][3] = 3;
    zoomprofile_2[i][4] = 1;
  }
  for(unsigned int i=1;i<=2;++i)
  {
    zoomprofile_2[i][0] = 4;
    zoomprofile_2[i][1] = 1;
    zoomprofile_2[i][2] = 5;
    zoomprofile_2[i][3] = 2;
    zoomprofile_2[i][4] = 1;
  }
  for(unsigned int i=3;i<=3;++i)
  {
    zoomprofile_2[i][0] = 6;
    zoomprofile_2[i][1] = 1;
    zoomprofile_2[i][2] = 3;
    zoomprofile_2[i][3] = 1;
    zoomprofile_2[i][4] = 1;
  }
  for(unsigned int i=4;i<=6;++i)
  {
    zoomprofile_2[i][0] = 3;
    zoomprofile_2[i][1] = 1;
    zoomprofile_2[i][2] = 1;
    zoomprofile_2[i][3] = 1;
    zoomprofile_2[i][4] = 1;
  }
  HelixRange top_range_2(0., 2.*M_PI,   -0.6, 0.6,   0.0, rho_max,   -0.9, 0.9,   -0.6, 0.6);
  sPHENIXTracker tracker_seeded(zoomprofile_2, 1, top_range_2, material, radii, Bfield);
  unsigned int max_hits = 500;
  unsigned int min_hits = 4;
  tracker_seeded.setNLayers(9);
  tracker_seeded.setSeedLayer(5);
  tracker_seeded.setRejectGhosts(true);
  tracker_seeded.setChi2Cut(8.0);
  tracker_seeded.setPrintTimings(true);
  tracker_seeded.setVerbosity(1);
  tracker_seeded.setCutOnDca(false);
  tracker_seeded.setSmoothBack(false);
  tracker_seeded.setBinScale(1.);
  tracker_seeded.setZBinScale(1.);
  tracker_seeded.setSeparateByHelicity(true);
  tracker_seeded.setRemoveHits(true);
  tracker_seeded.setChi2RemovalCut(3.0);
  
  vector<vector<unsigned int> > zoomprofile_3;
  zoomprofile_3.assign(7,onezoom);
  for(unsigned int i=0;i<=0;++i)
  {
    zoomprofile_3[i][0] = 10;
    zoomprofile_3[i][1] = 1;
    zoomprofile_3[i][2] = 5;
    zoomprofile_3[i][3] = 4;
    zoomprofile_3[i][4] = 1;
  }
  for(unsigned int i=1;i<=2;++i)
  {
    zoomprofile_3[i][0] = 8;
    zoomprofile_3[i][1] = 1;
    zoomprofile_3[i][2] = 2;
    zoomprofile_3[i][3] = 3;
    zoomprofile_3[i][4] = 1;
  }
  for(unsigned int i=3;i<=3;++i)
  {
    zoomprofile_3[i][0] = 6;
    zoomprofile_3[i][1] = 3;
    zoomprofile_3[i][2] = 1;
    zoomprofile_3[i][3] = 1;
    zoomprofile_3[i][4] = 1;
  }
  for(unsigned int i=4;i<=6;++i)
  {
    zoomprofile_3[i][0] = 3;
    zoomprofile_3[i][1] = 1;
    zoomprofile_3[i][2] = 1;
    zoomprofile_3[i][3] = 1;
    zoomprofile_3[i][4] = 1;
  }
  
  rho_max = 0.015;
  HelixRange top_range_3(0., 2.*M_PI,   -0.6, 0.6,   0.0, rho_max,   -0.9, 0.9,   -0.6, 0.6);
  sPHENIXTracker tracker_seeded_2(zoomprofile_3, 1, top_range_3, material, radii, Bfield);
  unsigned int max_hits_2 = 100;
  unsigned int min_hits_2 = 4;
  tracker_seeded_2.setNLayers(13);
  tracker_seeded_2.setSeedLayer(9);
  tracker_seeded_2.setRejectGhosts(true);
  tracker_seeded_2.setChi2Cut(8.0);
  tracker_seeded_2.setPrintTimings(true);
  tracker_seeded_2.setVerbosity(1);
  tracker_seeded_2.setCutOnDca(false);
  tracker_seeded_2.setSmoothBack(true);
  tracker_seeded_2.setBinScale(1.0);
  tracker_seeded_2.setZBinScale(1.0);
  tracker_seeded_2.setSeparateByHelicity(true);
  tracker_seeded_2.setRemoveHits(true);
  tracker_seeded_2.setChi2RemovalCut(3.0);
  
  
  TFile* mcfile = new TFile(argv[2], "recreate");
  TTree* mc_tree = new TTree("mc_events", "a tree of SimpleMCEvent");
  SimpleMCEvent* mcevent = new SimpleMCEvent();
  mc_tree->Branch("tracks", "SimpleMCEvent", &mcevent);
  TTree* reco_tree = new TTree("reco_events", "a tree of SimpleRecoEvent");
  SimpleRecoEvent* recoevent = new SimpleRecoEvent();
  reco_tree->Branch("tracks", "SimpleRecoEvent", &recoevent);
  
  for(unsigned int ev=0;ev<etree->GetEntries();ev++)
  {
    mcevent->tracks.clear();
    recoevent->tracks.clear();
    
    infile.cd();
    vector<float> pt_vec;
    
    etree->GetEntry(ev);
    cout<<"event "<<ev<<":"<<endl<<endl;
    
    vector<SimpleHit3D> hits_0_4;
    vector<SimpleHit3D> hits_5_8;
    vector<SimpleHit3D> hits_9_12;
    vector<SimpleHit3D> hits_all;
    vector<SimpleTrack3D> tracks_seed;
    vector<SimpleTrack3D> tracks;
    vector<SimpleTrack3D> tracks_2;
    
    vector<SimpleMCTrack>& mctracks = inputevent->tracks;
    
    cout<<"total MC tracks = "<<mctracks.size()<<endl;
    cout<<endl;
    
    // the last track contains noise hits
    unsigned int nfull=0;
    for(unsigned int trk=0;trk<(mctracks.size() - 1);trk++)
    {
      pt_vec.push_back(0.003*2./mctracks[trk].kappa);
      if(print_truth==true)
      {
        cout<<"truth track : "<<trk<<endl;
        cout<<"phi = "<<mctracks[trk].phi<<endl;
        cout<<"d = "<<mctracks[trk].d<<endl;
        cout<<"kappa = "<<mctracks[trk].kappa<<endl;
        cout<<"dzdl = "<<mctracks[trk].dzdl<<endl;
        cout<<"z0 = "<<mctracks[trk].z0<<endl;
        cout<<endl;
      }
      mcevent->tracks.push_back(SimpleMCTrack());
      mcevent->tracks.back().hits.clear();
      mcevent->tracks.back().kappa = mctracks[trk].kappa;
      mcevent->tracks.back().dzdl = mctracks[trk].dzdl;
      mcevent->tracks.back().d = mctracks[trk].d;
      mcevent->tracks.back().phi = mctracks[trk].phi;
      mcevent->tracks.back().z0 = mctracks[trk].z0;
      
      vector<SimpleMCHit>& mchits = mctracks[trk].hits;
      if(mchits.size() == 13){nfull+=1;}
      for(unsigned int hit=0;hit<mchits.size();hit++)
      {
        SimpleMCHit& mchit = mchits[hit];
        
        float phi = atan2(mchit.y, mchit.x);
        float xy_error = smear_xy_layer[mchit.layer]*0.5;
        float x_error = fabs(xy_error*sin(phi));
        float y_error = fabs(xy_error*cos(phi));
        float z_error = smear_z_layer[mchit.layer]*0.5;
        
        mcevent->tracks.back().hits.push_back(SimpleMCHit(mchit.x, mchit.y, mchit.z, mchit.layer, mchit.index));
        
        if(mchit.layer<5)
        {
          hits_0_4.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        }
        else if(mchit.layer<9)
        {
          hits_5_8.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        }
        else
        {
          hits_9_12.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        }
        hits_all.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
        if(print_truth==true)
        {
          cout<<mchit.x<<" "<<mchit.y<<" "<<mchit.z<<" "<<mchit.index<<endl;
        }
      }
      if(print_truth==true){cout<<endl<<endl;}
    }
    cout<<"nfull = "<<nfull<<endl;
    mcfile->cd();
    mc_tree->Fill();
    
    vector<SimpleMCHit>& mchits = mctracks.back().hits;
    for(unsigned int h=0;h<mchits.size();h++)
    {
      SimpleMCHit& mchit = mchits[h];
      
      float phi = atan2(mchit.y, mchit.x);
      float xy_error = smear_xy_layer[mchit.layer]*0.5;
      float x_error = fabs(xy_error*sin(phi));
      float y_error = fabs(xy_error*cos(phi));
      float z_error = smear_z_layer[mchit.layer]*0.5;
      
      if(mchit.layer<5)
      {
        hits_0_4.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
      }
      else if(mchit.layer<9)
      {
        hits_5_8.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
      }
      else
      {
        hits_9_12.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
      }
      hits_all.push_back(SimpleHit3D(mchit.x,x_error, mchit.y,y_error, mchit.z,z_error, mchit.index, mchit.layer));
    }
    
    
    timeval t1,t2;
    double time1=0.;
    double time2=0.;
    unsigned int ngood = 0;
    unsigned int n_trk = tracks.size();

    tracks.clear();
    tracks_seed.clear();
    tracker.clear();
    tracker_seeded.clear();
    tracker_seeded_2.clear();
    
    gettimeofday(&t1, NULL);
//     tracker.findHelices(hits_0_4, min_hits_seed, max_hits_seed, tracks_seed);
    tracker.findHelicesParallel(hits_0_4, min_hits_seed, max_hits_seed, tracks_seed);
    gettimeofday(&t2, NULL);
    time1 = ((double)(t1.tv_sec) + (double)(t1.tv_usec)/1000000.);
    time2 = ((double)(t2.tv_sec) + (double)(t2.tv_usec)/1000000.);
    cout<<"seed tracking time = "<<(time2 - time1)<<endl<<endl;
    tracker_seeded.setSeedStates(tracker.getKalmanStates());
    
    cout<<tracks_seed.size()<<" track seeds found"<<endl<<endl;
    
//     tracks = tracks_seed;
    
    gettimeofday(&t1, NULL);
    tracker_seeded.findSeededHelices(tracks_seed, hits_5_8, min_hits, max_hits, tracks);
    gettimeofday(&t2, NULL);
    time1 = ((double)(t1.tv_sec) + (double)(t1.tv_usec)/1000000.);
    time2 = ((double)(t2.tv_sec) + (double)(t2.tv_usec)/1000000.);
    cout<<"tracking time stage 1 = "<<(time2 - time1)<<endl<<endl;
//     vector<float>& isolation = tracker_seeded.getIsolation();
    tracker_seeded_2.setSeedStates(tracker_seeded.getKalmanStates());
    
    cout<<tracks.size()<<" stage 1 tracks found"<<endl<<endl;
    
    gettimeofday(&t1, NULL);
    tracker_seeded_2.findSeededHelices(tracks, hits_9_12, min_hits_2, max_hits_2, tracks_2);
    gettimeofday(&t2, NULL);
    time1 = ((double)(t1.tv_sec) + (double)(t1.tv_usec)/1000000.);
    time2 = ((double)(t2.tv_sec) + (double)(t2.tv_usec)/1000000.);
    cout<<"tracking time stage 2 = "<<(time2 - time1)<<endl<<endl;
    vector<float>& isolation = tracker_seeded_2.getIsolation();
    
    for(unsigned int trk=0;trk<tracks.size();++trk)
    {
      if(tracker_seeded_2.seedWasUsed(trk) == false)
      {
        tracks_2.push_back(tracks[trk]);
        tracker_seeded_2.getKalmanStates().push_back(tracker_seeded.getKalmanStates()[trk]);
        isolation.push_back(tracker.getIsolation()[trk]);
      }
    }
    
    cout<<tracks_2.size()<<" stage 2 tracks found"<<endl<<endl;
    
    
    
    n_trk = tracks_2.size();
    cout<<endl;
    
    
    if(print_reco==true)
    {
      for(unsigned int trk=0;trk<n_trk;++trk)
      {
        cout<<"track "<<trk<<" : "<<endl;
        unsigned int n_hit = tracks_2[trk].hits.size();
        for(unsigned int ht=0;ht<n_hit;++ht)
        {
          cout<<"hit "<<ht<<" : ";
          cout<<tracks_2[trk].hits[ht].x<<" ";
          cout<<tracks_2[trk].hits[ht].y<<" ";
          cout<<tracks_2[trk].hits[ht].z<<" ";
          cout<<tracks_2[trk].hits[ht].index<<endl;
        }
        cout<<"phi = "<<tracks_2[trk].phi<<endl;
        cout<<"d = "<<tracks_2[trk].d<<endl;
        cout<<"kappa = "<<tracks_2[trk].kappa<<endl;
        cout<<"dzdl = "<<tracks_2[trk].dzdl<<endl;
        cout<<"z0 = "<<tracks_2[trk].z0<<endl;
        cout<<endl;
      }
    }
    
    
    
    for(unsigned int trk=0;trk<n_trk;++trk)
    {
      recoevent->tracks.push_back(SimpleRecoTrack());
      recoevent->tracks.back().kappa = tracks_2[trk].kappa;
      recoevent->tracks.back().phi = tracks_2[trk].phi;
      recoevent->tracks.back().d = tracks_2[trk].d;
      recoevent->tracks.back().z0 = tracks_2[trk].z0;
      recoevent->tracks.back().dzdl = tracks_2[trk].dzdl;
      for(unsigned int h=0;h<tracks_2[trk].hits.size();++h)
      {
        recoevent->tracks.back().indexes.push_back(tracks_2[trk].hits[h].index);
      }
      recoevent->tracks.back().chi2 = (tracker.getKalmanStates())[trk].chi2/(2.*13.-5.);
      recoevent->tracks.back().isolation = isolation[trk];
    }
    mcfile->cd();
    reco_tree->Fill();
    
  }
  
  
  mcfile->cd();
  mc_tree->Write();
  reco_tree->Write();
  mcfile->Close();
  mcfile->Delete();
  
  return 0;
}


